use std::rc::Rc;
use rand::{Rng, thread_rng};
use rand::distributions::Uniform;
use twinkle::datasets::load_digits::load_digits;
use twinkle::nn::functional as F;
use twinkle::Tensor;

pub fn fit() {
    let (x_train, x_test, y_train, y_test) = load_digits();

    let epochs = 50u32;
    let lr = Tensor::scalar(1e-4f64).expand(&[64, 10]);

    let alpha = Tensor::scalar(1e-3f64).expand(&[64, 10]);
    let decay = &Tensor::scalar(1f64).expand(&[64, 10]) - &(&lr * &alpha);

    let layer1_data = thread_rng()
        .sample_iter(Uniform::new(-0.01, 0.01))
        .take(64 * 32)
        .collect::<Box<[f64]>>();
    let mut layer1: Tensor<f64> = (Rc::new(layer1_data), [64, 32].into()).into();

    let layer2_data = thread_rng()
        .sample_iter(Uniform::new(-0.01, 0.01))
        .take(32 * 10)
        .collect::<Box<[f64]>>();
    let mut layer2: Tensor<f64> = (Rc::new(layer2_data), [32, 10].into()).into();

    let mut train_score = 0u32;
    let mut test_score = 0u32;

    for epoch in 0..epochs {
        for i in 0..x_train.shape()[0] {
            let x_i = x_train.index(&[i]).reshape(&[1, 64]);
            let y_i = y_train.index(&[i]).reshape(&[1, 10]);

            let layer1_out = F::relu(&x_i.mm(&layer1));
            let y_hat = layer1_out.mm(&layer2);

            if y_hat.argmax().item() == y_i.argmax().item() {
                train_score += 1;
            }

            let layer2_err = &y_hat - &y_i;
            let layer1_err = layer2_err.mm(&layer2.t());

            layer2 = &(&layer2 * &decay) - &(&lr * &(&layer1_out * &F::relu_grad(&layer1_out)).t()
                .mm(&layer2_err));
            layer1 = &(&layer1 * &decay) - &(&lr * &x_i.t()
                .mm(&layer1_err));
        }

        for i in 0..x_test.shape()[0] {
            let x_i = x_test.index(&[i]).reshape(&[1, 64]);
            let y_i = y_test.index(&[i]).reshape(&[1, 10]);

            let layer1_out = F::relu(&x_i.mm(&layer1));
            let y_hat = layer1_out.mm(&layer2);

            if y_hat.argmax().item() == y_i.argmax().item() {
                test_score += 1;
            }
        }

        print!("epoch {:2}, ", epoch);
        print!("train accuracy {:.3}, ", f64::from(train_score) / 1000f64);
        println!("test accuracy {:.3}", f64::from(test_score) / 797f64);

        train_score = 0;
        test_score = 0;
    }
}

fn main() {
    fit();
}
